import os, numpy as np, argparse, json, sys, numba, yaml, shutil
import multiprocessing

from tensorflow import atan2

# import torch.multiprocessing as multiprocessing
import mot_3d.visualization as visualization, mot_3d.utils as utils
from mot_3d.data_protos import BBox, Validity, SearchBBox
from mot_3d.mot import MOTModel
from mot_3d.frame_data import FrameData
from data_loader import WaymoLoader
from ipdb import set_trace
from mot_3d.utils import Timer
import time
from cyw_devkit import CywDataset

timer = Timer(10)

class TicToc():
    def __init__(self, msgs: str):
        self.start_time = time.time()  # 程序开始时间
        self.end_time = time.time()  # 程序开始时间
        self.msgs = msgs

    def tic(self):
        self.start_time = time.time()  # 程序开始时间

    def toc(self):
        self.end_time = time.time()  # 程序结束时间
        run_time = self.end_time - self.start_time  # 程序的运行时间，单位为秒
        print("-----------  <{}> cost [{}] ms   ---------------".format(self.msgs, run_time * 1000))


parser = argparse.ArgumentParser()
# running configurations
parser.add_argument('--config_path', type=str, default='configs/cyw_configs/immortal_for_eog_drive.yaml')
parser.add_argument('--data_path', type=str, default='/media/adt/T7/ZWH/docker/files/data/motion/datasets')
args = parser.parse_args()

class2idx = {
    'Car': 0,
    'Pedestrian': 1,
    'Bus': 2,
    'Truck': 3,
}
idx2class = {val: key for key, val in class2idx.items()}



def load_gt_bboxes(gt_folder, data_folder, segment_name, type_token):
    gt_info = np.load(os.path.join(gt_folder, '{:}.npz'.format(segment_name)),
                      allow_pickle=True)
    ego_info = np.load(os.path.join(data_folder, 'ego_info', '{:}.npz'.format(segment_name)),
                       allow_pickle=True)
    bboxes, ids, inst_types = gt_info['bboxes'], gt_info['ids'], gt_info['types']
    gt_ids, gt_bboxes = utils.inst_filter(ids, bboxes, inst_types, type_field=[type_token], id_trans=True)

    ego_keys = sorted(utils.str2int(ego_info.keys()))
    egos = [ego_info[str(key)] for key in ego_keys]
    gt_bboxes = gt_bbox2world(gt_bboxes, egos)
    return gt_bboxes, gt_ids


def gt_bbox2world(bboxes, egos):
    frame_num = len(egos)
    for i in range(frame_num):
        ego = egos[i]
        bbox_num = len(bboxes[i])
        for j in range(bbox_num):
            bboxes[i][j] = BBox.bbox2world(ego, bboxes[i][j])
    return bboxes


def frame_visualization(bboxes, ids, states, gt_bboxes=None, gt_ids=None, pc=None, dets=None, name=''):
    visualizer = visualization.Visualizer2D(name=name, figsize=(12, 12))
    if pc is not None:
        visualizer.handler_pc(pc)
    for _, bbox in enumerate(gt_bboxes):
        visualizer.handler_box(bbox, message='', color='black')
    dets = [d for d in dets if d.s >= 0.1]
    for det in dets:
        visualizer.handler_box(det, message='%.2f' % det.s, color='green', linestyle='dashed')
    for _, (bbox, id, state) in enumerate(zip(bboxes, ids, states)):
        if Validity.agein1(state):
            visualizer.handler_box(bbox, message=str(id), color='red')
        else:
            visualizer.handler_box(bbox, message=str(id), color='light_blue')
    visualizer.show()
    visualizer.save('temp.jpg')
    visualizer.close()
    import pdb
    pdb.set_trace()


def sequence_mot(configs, data_loader: WaymoLoader, return_tracklets=False):
    tracker = MOTModel(configs)
    frame_num = len(data_loader)
    IDs, bboxes, states, types, save_paths, stamps, b2l_tfs = list(), list(), list(), list(), list(), list(), list()
    for idx, data in enumerate(data_loader):
        if data['label'] is None:
            continue

        stamp = data['label'].header.stamp
        ego = data_loader.data_dict['super_transform'].localizer.get_tf(stamp=stamp)
        det_ids = []
        det_types = []
        dets = []
        dummy_states = []
        for obj_ids, obj_type, tensor in zip(data['label'].data.obj_ids, data['label'].data.obj_types, data['label'].data.tensors):
            # if obj_type in class2idx.keys():
            det_ids.append(obj_ids)
            det_types.append(obj_type)
            dets.append(
                np.array((tensor[0], tensor[1], tensor[2], tensor[6], tensor[3], tensor[4], tensor[5], 1.0)))
            dummy_states.append('alive_1_0')

        # det_locals = []
        # for det in dets:
        #     det_local = BBox.bbox2world(ego, det)
        #     det_locals.append(det_local)

        frame_data = {
            'time_stamp_int': stamp,
            'time_stamp': stamp / 1000.0,
            'dets': dets,
            'ego': ego,
            'pc': None,
            'det_types': det_types,
            'aux_info': {
                'is_key_frame': True,
                'velos': None
            },
            'abs_frame_index': data_loader.keyframe_idxs.index(idx),
        }
        frame_data = FrameData(dets=frame_data['dets'], ego=frame_data['ego'], pc=frame_data['pc'],
                               det_types=frame_data['det_types'], aux_info=frame_data['aux_info'],
                               time_stamp=frame_data['time_stamp'], abs_frame_index=frame_data['abs_frame_index'])

        # # mot
        # frame_results = tracker.frame_mot(frame_data)
        # result_pred_bboxes = [trk['bboxes'] for trk in frame_results]
        # result_pred_ids = [trk['id'] for trk in frame_results]
        # result_pred_states = [trk['state'] for trk in frame_results]
        # result_types = [trk['type'] for trk in frame_results]

        for i in range(len(frame_data.dets)):
            frame_data.dets[i] = BBox.bbox2world(ego, frame_data.dets[i])

        # zwh: save to label_tracker
        save_path = os.path.join(data_loader.father_path,
                                 data_loader.data_dict['frames'][idx]['label'].replace('label', 'label_e2e'))


        IDs.append(det_ids)
        bboxes.append(frame_data.dets)
        states.append(dummy_states)
        types.append(det_types)
        save_paths.append(save_path)
        stamps.append(stamp)
        b2l_tfs.append(ego)

        # # wrap for output
        # IDs.append(result_pred_ids)
        # # result_pred_bboxes = [BBox.bbox2array(bbox) for bbox in result_pred_bboxes]
        # bboxes.append(result_pred_bboxes)
        # states.append(result_pred_states)
        # types.append(result_types)
        # save_paths.append(save_path)
        # stamps.append(stamp)
        # b2l_tfs.append(ego)
    if return_tracklets:
        return tracker.tracklets
    return IDs, bboxes, states, types, save_paths, stamps, b2l_tfs

def transform_box2_to_box1(x1, y1, yaw1, x2, y2, yaw2, without_yaw=False):
    # 计算相对位移
    dx = x2 - x1
    dy = y2 - y1
    if without_yaw:
        return dx, dy, 0

    # 坐标旋转
    cos_yaw1 = np.cos(yaw1)
    sin_yaw1 = np.sin(yaw1)
    x_in_box1 = dx * cos_yaw1 + dy * sin_yaw1
    y_in_box1 = -dx * sin_yaw1 + dy * cos_yaw1

    # 角度转换与归一化
    yaw_in_box1 = yaw2 - yaw1
    yaw_in_box1 = (yaw_in_box1 + np.pi) % (2 * np.pi) - np.pi  # 归一化到[-π, π)

    return x_in_box1, y_in_box1, yaw_in_box1


def inverse_cumsum(cumsum_array):
    original = [cumsum_array[0]]  # 处理第一个元素
    for i in range(1, len(cumsum_array)):
        original.append(cumsum_array[i] - cumsum_array[i-1])
    return np.stack(original,axis=0)


def main_worker(dataset_path, args):
    laber_e2e_dir = os.path.join(dataset_path, 'samples/label_e2e')
    if not os.path.exists(laber_e2e_dir):
        os.makedirs(laber_e2e_dir)
    if "town" in dataset_path or "carla" in dataset_path:
        radar_dim = 7
    else:
        radar_dim = 5
    data_loader = CywDataset(
        dataset_path=dataset_path, radar_dim=radar_dim)

    ids, bboxes, states, types, save_paths, stamps, b2l_tfs = sequence_mot(args.configs, data_loader)

    # 同一id使用最后一帧尺寸
    const_bboxes = []
    boxes_states = dict()
    for id, bbox, type in zip(ids[::-1], bboxes[::-1], types[::-1]):
        const_bbox = []
        for single_id, single_box, single_type in zip(id, bbox, type):
            if single_id not in boxes_states:
                boxes_states[single_id] = single_box
            else:
                single_box.w = boxes_states[single_id].w
                single_box.l = boxes_states[single_id].l
                single_box.h = boxes_states[single_id].h
            if single_type == 'Pedestrian':  # 行人yaw方向改成速度方向
                single_box.o = float(atan2(single_box.vy, single_box.vx))
            const_bbox.append(single_box)
        const_bboxes.append(const_bbox)
    const_bboxes = const_bboxes[::-1]

    assert len(bboxes) == len(save_paths)

    search_bbox = SearchBBox()
    for result_pred_ids, stamp, result_pred_bboxes in zip(ids, stamps, const_bboxes):
        for id, box in zip(result_pred_ids, result_pred_bboxes):
            search_bbox.insert(id, box, stamp)

    step = 1000 * 0.5  # cyw数据命名为时间ms
    num = 12  # 轨迹点数量
    ego_num_pre = 12  # 自车历史轨迹点数量
    ego_num_post = 6  # 自车未来轨迹点数量
    for result_pred_ids, result_types, result_pred_bboxes, result_pred_states, save_path, stamp, b2l_tf in zip(ids,
                                                                                                               types,
                                                                                                               const_bboxes,
                                                                                                               states,
                                                                                                               save_paths,
                                                                                                               stamps,
                                                                                                               b2l_tfs):

        # planning
        cur_pose = data_loader.data_dict['super_transform'].localizer.get_tf(stamp)
        ego_fut_trajs_pre = []
        for i in range(ego_num_pre):
            stamp_pre = stamp - (i + 1) * step
            if stamp_pre < data_loader.data_dict['super_transform'].localizer.time_start:
                stamp_pre = data_loader.data_dict['super_transform'].localizer.time_start
            pre_pose = data_loader.data_dict['super_transform'].localizer.get_tf(stamp_pre)  # cyw内部是ms
            pre_xyzi = pre_pose[:, 3]
            pre_xyzi_cur = np.linalg.inv(cur_pose) @ pre_xyzi
            ego_fut_trajs_pre.append(pre_xyzi_cur[:2])
        ego_fut_trajs_pre = np.stack(ego_fut_trajs_pre, 0)

        ego_fut_trajs = []
        for i in range(ego_num_post):
            stamp_post = stamp + (i + 1) * step
            if stamp_post > data_loader.data_dict['super_transform'].localizer.time_end:
                stamp_post = data_loader.data_dict['super_transform'].localizer.time_end
            post_pose = data_loader.data_dict['super_transform'].localizer.get_tf(stamp_post)
            post_xyzi = post_pose[:, 3]
            post_xyzi_cur = np.linalg.inv(cur_pose) @ post_xyzi
            ego_fut_trajs.append(post_xyzi_cur[:2])
        ego_fut_trajs = np.stack(ego_fut_trajs, 0)

        # drive command according to final fut step offset
        if ego_fut_trajs[-1][1] <= -2:
            command = np.array([0, 0, 1])  # Turn Right，右转
        elif ego_fut_trajs[-1][1] >= 2:
            command = np.array([1, 0, 0])  # Turn Left，左转
        else:
            command = np.array([0, 1, 0])  # Go Straight，直行

        # 轨迹全部换成增量
        ego_fut_trajs_pre = inverse_cumsum(ego_fut_trajs_pre)
        ego_fut_trajs = inverse_cumsum(ego_fut_trajs)

        json_data = dict()
        data = {
            'gt_ego_fut_trajs': {
                'x': ego_fut_trajs[:, 0].tolist(),
                'y': ego_fut_trajs[:, 1].tolist(),
            },
            'gt_ego_fut_trajs_pre': {
                'x': ego_fut_trajs_pre[:, 0].tolist(),
                'y': ego_fut_trajs_pre[:, 1].tolist(),
            },
            'gt_ego_fut_cmd': command.tolist(),
        }
        json_data['ego'] = data
        json_data['obj'] = []
        for id, obj_type, cur_box, state_string in zip(result_pred_ids, result_types, result_pred_bboxes,
                                                       result_pred_states):
            tokens = state_string.split('_')
            if len(tokens) < 3:
                continue
            if tokens[0] == 'alive' and int(tokens[1]) == 0:
                continue

            cur_box_2_base = BBox.bbox2world(np.linalg.inv(b2l_tf), cur_box)

            # 目标速度，因为carla数据不经过跟踪
            fut_stamp = stamp + 1000 #1s后位置计算速度
            fut_box = search_bbox.search(id, fut_stamp)
            vel = np.array([[0.0,0.0,0.0]])
            vel[0,0] = (fut_box.x-cur_box.x)/1.0
            vel[0,1] = (fut_box.y-cur_box.y)/1.0
            vel_new = (np.linalg.inv(b2l_tf[:3,:3]) @vel.T).T

            # 未来轨迹
            fut_trajs = np.zeros((num, 2))
            for idx in range(num):
                fut_stamp = stamp + (idx + 1) * step
                fut_box = search_bbox.search(id, fut_stamp)

                fut_box_2_base = BBox.bbox2world(np.linalg.inv(b2l_tf), fut_box)

                # x_t, y_t, yaw_t = transform_box2_to_box1(cur_box_2_base.x, cur_box_2_base.y, cur_box_2_base.o,
                #                                          fut_box_2_base.x, fut_box_2_base.y, fut_box_2_base.o)
                # 不考虑目标坐标系的yaw角，参考sparsedrive，免得恢复角度时候受到yaw误差的干扰。似乎drive transformer还是执意考虑yaw
                x_t, y_t, yaw_t = transform_box2_to_box1(cur_box_2_base.x, cur_box_2_base.y, cur_box_2_base.o,
                                                         fut_box_2_base.x, fut_box_2_base.y, fut_box_2_base.o,
                                                         without_yaw=True)

                fut_trajs[idx][0] = x_t
                fut_trajs[idx][1] = y_t
                pass
            fut_trajs2 = inverse_cumsum(fut_trajs)

            data = {
                'obj_id': str(id),
                'obj_type': obj_type,
                'psr': {
                    'position': {
                        'x': cur_box_2_base.x,
                        'y': cur_box_2_base.y,
                        'z': cur_box_2_base.z
                    },
                    'rotation': {
                        'x': 0,
                        'y': 0,
                        'z': cur_box_2_base.o
                    },
                    'scale': {
                        'x': cur_box_2_base.l,
                        'y': cur_box_2_base.w,
                        'z': cur_box_2_base.h
                    },
                    'velocity': {
                        'x': vel_new[0,0],
                        'y': vel_new[0,1],
                        'z': vel_new[0,2]
                    },
                    'gt_agent_fut_trajs': {
                        'x': fut_trajs2[:, 0].tolist(),
                        'y': fut_trajs2[:, 1].tolist(),
                    },
                }
            }
            json_data['obj'].append(data)
        with open(save_path, 'w', encoding='utf-8') as f:
            # print(save_path)
            json.dump(json_data['obj'], f, ensure_ascii=False, indent=2)
        if 'ego' in json_data:
            with open(save_path.replace('.json', '-ego.json'), 'w', encoding='utf-8') as f:
                json.dump(json_data['ego'], f, ensure_ascii=False, indent=2)

    pass

def main(args):
    # 统计耗时
    cost = TicToc("e2e label 生成")
    assert os.path.exists(args.data_path)
    datasets = os.listdir(args.data_path)
    datasets.sort(key=lambda x: x)
    files = []
    for dataset in datasets:
        if "__" == dataset[:2]:  # 数据集以"__"开头
            dataset_path = os.path.join(args.data_path, dataset)
            # assert check_dataset_is_splited(dataset_path)  # 上一步已经完成
            files.append(dataset_path)

    # load model configs
    configs = yaml.safe_load(open(args.config_path, 'r'))
    args.configs = configs

    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


if __name__ == '__main__':
    main(args)


